from io import BytesIO
import numpy as np
import cv2 as cv
import matplotlib.pyplot as plt
import base64
import re

class ImageUtil:

    @staticmethod
    def from_array(source):
        source = base64.b64decode(re.sub('^data:image/.+;base64,', '', source))
        return cv.imdecode(np.frombuffer(source, np.uint8), cv.IMREAD_COLOR)

    @staticmethod
    def to_bytes(source):
        return base64.b64decode(re.sub('^data:image/.+;base64,', '', source))
    
    @staticmethod
    def to_base64(img):
        if type(img) == bytes:
            return base64.b64encode(img).decode('utf-8')
        iobytes = cv.imencode(".jpg", img)[1].tobytes()
        return base64.b64encode(iobytes).decode('utf-8')


        
        

    @staticmethod
    def grayfy(img):
        "图像灰度化"
        return cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    
    @staticmethod
    def binarify(img, threshold=100, maxval=200, **args):
        "二值化"
        ret, binary = cv.threshold(img, threshold, maxval, cv.THRESH_BINARY)
        return binary
    
    @staticmethod
    def rotate(img, theta=90, **args):
        "旋转"
        rows, cols, c = img.shape
        M = cv.getRotationMatrix2D(((cols - 1) / 2.0, (rows - 1) / 2.0), theta, 1)
        return cv.warpAffine(img, M, (cols, rows))

    @staticmethod
    def scale(img, fx=5, fy=5, **args):
        "放缩"
        height, width, channel = img.shape
        return cv.resize(img, (fx * width, fy * height), interpolation=cv.INTER_CUBIC)

    @staticmethod
    def move(img, fx, fy, **args):
        "平移"
        height, width, channel = img.shape
        M = np.float32([[1, 0, fy], [0, 1, fx]])
        return cv.warpAffine(img, M, (width, height))

    @staticmethod
    def dft(img):
        "傅里叶变换"
        img = ImageUtil.grayfy(img)
        rows, cols = img.shape
        nrows = cv.getOptimalDFTSize(rows)
        ncols = cv.getOptimalDFTSize(cols)
        nimg = np.zeros((nrows,ncols))
        nimg[:rows,:cols] = img
        img = nimg
        dft = cv.dft(np.float32(img), flags=cv.DFT_COMPLEX_OUTPUT)
        dft_shift = np.fft.fftshift(dft)
        magnitude_spectrum = 20 * np.log(cv.magnitude(dft_shift[:, :, 0], dft_shift[:, :, 1]))
        # plt.subplot(121),plt.imshow(img, cmap = 'gray')
        # plt.title('Input Image'), plt.xticks([]), plt.yticks([])
        # plt.subplot(122),plt.imshow(magnitude_spectrum, cmap = 'gray')
        # plt.title('Magnitude Spectrum'), plt.xticks([]), plt.yticks([])
        # plt.show()
        return magnitude_spectrum

    @staticmethod
    def hist(img):
        "直方图"
        hist_b = cv.calcHist([img], [0], None, [256], [0, 256])
        hist_g = cv.calcHist([img], [1], None, [256], [0, 256])
        hist_r = cv.calcHist([img], [2], None, [256], [0, 256])
        plt.figure() 
        plt.plot(hist_b, 'b')
        plt.plot(hist_g, 'g')
        plt.plot(hist_r, 'r')
        plt.xlim([0, 256])
        res = BytesIO()
        plt.savefig(res)
        return res.getvalue()
        # plt.show()
    
    @staticmethod
    def hist_equal(img):
        "直方图均衡化"
        b = cv.equalizeHist(img[:,:,0])
        g = cv.equalizeHist(img[:,:,1])
        r = cv.equalizeHist(img[:,:,2])
        res = np.zeros_like(img)
        res[:,:,0] = b
        res[:,:,1] = g
        res[:,:,2] = r
        return res
    
    @staticmethod
    def adaptive_threshold(img):
        "自适应阈值"
        img = cv.medianBlur(img[:,:,0], 5) # 只提取一个维度
        th = cv.adaptiveThreshold(img,255,cv.ADAPTIVE_THRESH_GAUSSIAN_C,\
            cv.THRESH_BINARY,11,2)
        return th
    
    @staticmethod
    def blurring(img, kind='gaussian', **args):
        "图像平滑"
        if kind == 'gaussian':
            return cv.GaussianBlur(img, (5, 5), 0)
        elif kind == 'avg':
            return cv.blur(img, (5, 5))
        elif kind == 'median':
            return cv.medianBlur(img, 5)
        elif kind == 'bilateral':
            return cv.bilateralFilter(img, 9, 75, 75)
    
    @staticmethod
    def morphology(img, kind='erosion', **args):
        "形态学变换"
        kernel = np.ones((5,5),np.uint8)
        if kind == 'erosion':
            return cv.erode(img,kernel,iterations = 1)
        elif kind == 'dilation':
            return cv.dilate(img, kernel, iterations=1)
        elif kind == 'opening':
            return cv.morphologyEx(img, cv.MORPH_OPEN, kernel)
        elif kind == 'closing':
            return cv.morphologyEx(img, cv.MORPH_CLOSE, kernel)
        elif kind == 'gradient':
            return cv.morphologyEx(img, cv.MORPH_GRADIENT, kernel)
        elif kind == 'tophat':
            return cv.morphologyEx(img, cv.MORPH_TOPHAT, kernel)
        elif kind == 'blackhat':
            return cv.morphologyEx(img, cv.MORPH_BLACKHAT, kernel)
        
    @staticmethod
    def gradient(img, kind='laplacian', **args):
        "图像梯度"
        if kind == 'laplacian':
            return cv.Laplacian(img,cv.CV_64F)
        elif kind == 'sobelx':
            return cv.Sobel(img,cv.CV_64F,1,0,ksize=5)
        elif kind == 'sobely':
            return cv.Sobel(img,cv.CV_64F,0,1,ksize=5)
    
    @staticmethod
    def canny_edge(img):
        "图像边缘"
        edges = cv.Canny(img,100,200)
        return edges

    @staticmethod
    def hough(img, kind='normal', **args):
        "霍夫变换"
        gray = cv.cvtColor(img,cv.COLOR_BGR2GRAY)
        edges = cv.Canny(gray,50,150,apertureSize = 3)
        if kind == 'p':
            lines = cv.HoughLinesP(edges,1,np.pi/180,100,minLineLength=100,maxLineGap=10)
            for line in lines:
                x1,y1,x2,y2 = line[0]
                cv.line(img,(x1,y1),(x2,y2),(0,255,0),2)
            return img
        elif kind == 'circle':
            img = cv.cvtColor(img,cv.COLOR_BGR2GRAY)
            img = cv.medianBlur(img,5)
            cimg = cv.cvtColor(img,cv.COLOR_GRAY2BGR)
            circles = cv.HoughCircles(img,cv.HOUGH_GRADIENT,1,20,
                            param1=50,param2=30,minRadius=0,maxRadius=0)
            circles = np.uint16(np.around(circles))
            for i in circles[0,:]:
                # draw the outer circle
                cv.circle(cimg,(i[0],i[1]),i[2],(0,255,0),2)
                # draw the center of the circle
                cv.circle(cimg,(i[0],i[1]),2,(0,0,255),3)
            # plt.imshow(img)
            # plt.show()
            return img
        else:
            lines = cv.HoughLines(edges,1,np.pi/180,200)
            for line in lines:
                rho,theta = line[0]
                a = np.cos(theta)
                b = np.sin(theta)
                x0 = a*rho
                y0 = b*rho
                x1 = int(x0 + 1000*(-b))
                y1 = int(y0 + 1000*(a))
                x2 = int(x0 - 1000*(-b))
                y2 = int(y0 - 1000*(a))
                cv.line(img,(x1,y1),(x2,y2),(0,0,255),2)
                return img
    
    @staticmethod
    def segmentation(img):
        gray = cv.cvtColor(img,cv.COLOR_BGR2GRAY)
        ret, thresh = cv.threshold(gray,0,255,cv.THRESH_BINARY_INV+cv.THRESH_OTSU)
        kernel = np.ones((3,3),np.uint8)
        opening = cv.morphologyEx(thresh,cv.MORPH_OPEN,kernel, iterations = 2)
        sure_bg = cv.dilate(opening,kernel,iterations=3)
        dist_transform = cv.distanceTransform(opening,cv.DIST_L2,5)
        ret, sure_fg = cv.threshold(dist_transform,0.7*dist_transform.max(),255,0)
        sure_fg = np.uint8(sure_fg)
        unknown = cv.subtract(sure_bg,sure_fg)
        ret, markers = cv.connectedComponents(sure_fg)
        markers = markers+1
        markers[unknown==255] = 0
        markers = cv.watershed(img,markers) #语义
        img[markers == -1] = [255,0,0] # 标识图
        # plt.imshow(markers)
        # plt.show()
        return img
    
    @staticmethod
    def denoising(img):
        "去噪"
        return cv.fastNlMeansDenoisingColored(img,None,10,10,7,21)








        
    

# if __name__ == '__main__':
#     # path = input('输入一个文件路径: ')
#     image = ImageUtil.read_image('./coins.jpg')
#     image = ImageUtil.denoising(image)
#     ImageUtil.show_image(image)
#     # ImageUtil.show_image(image)
